package com.learn.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.BroadcastConnectedStream;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.util.Collector;

import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * 需求：公司有10个广告位, 其广告的内容(描述和图片)会经常变动（广告到期，更换广告等）
 * 实现：
 * 1）通过socket输入广告id
 * 2）关联出来广告的信息打印出来，如果广告发生改变，能够感知到
 * 以一个自定义source来模拟操作过程
 */
public class BroadcastStateDemo {
    public static void main(String[] args) throws Exception {
        /**
         * 实现步骤；
         * 1）初始化flink流式处理的运行环境
         * 2）设置checkpoint周期性运行，每5秒钟运行一次
         * 3）构建数据源（两个数据源）
         * 4）将广告流作为配置流转换成广播流
         * 5）将两个流进行关联操作（connect）
         * 6）对关联的数据进行拉宽操作
         * 7）打印输出
         * 8）运行作业
         */

        //TODO 1）初始化flink流式处理的运行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        //TODO 2）设置checkpoint周期性运行，每5秒钟运行一次
        env.enableCheckpointing(5000L);

        //TODO 3）构建数据源（两个数据源）
        //构建事件流
        DataStreamSource<String> socketTextStream = env.socketTextStream("node1", 9999);
        SingleOutputStreamOperator<Integer> adIdDataStream = socketTextStream.map(new MapFunction<String, Integer>() {
            @Override
            public Integer map(String s) throws Exception {
                return Integer.parseInt(s);
            }
        });

        //构建广播流（广告位）
        DataStreamSource<Map<Integer, Tuple2<String, String>>> adSourceStream = env.addSource(new MySourceForBroadcastFunction());
        //将广告位流广播出去
        MapStateDescriptor<Integer, Tuple2<String, String>> mapStateDescriptor =
                new MapStateDescriptor<Integer, Tuple2<String, String>>(
                        "broadcaststate",
                        TypeInformation.of(new TypeHint<Integer>() { }),
                        TypeInformation.of(new TypeHint<Tuple2<String, String>>() { })
                );

        //TODO 4）将广告流作为配置流转换成广播流
        //广告位流变成了广播流
        BroadcastStream<Map<Integer, Tuple2<String, String>>> broadcastStream = adSourceStream.broadcast(mapStateDescriptor);

        //TODO 5）将两个流进行关联操作（connect）
        BroadcastConnectedStream<Integer, Map<Integer, Tuple2<String, String>>> connectedStream = adIdDataStream.connect(broadcastStream);

        //TODO 6）对关联的数据进行拉宽操作
        SingleOutputStreamOperator<Tuple2<String, String>> result = connectedStream.process(new BroadcastProcessFunction<Integer,
                        Map<Integer, Tuple2<String, String>>, Tuple2<String, String>>() {

            //定义state的描述器
            MapStateDescriptor<Integer, Tuple2<String, String>> mapStateDescriptor =
                    new MapStateDescriptor<Integer, Tuple2<String, String>>(
                            "broadcaststate",
                            TypeInformation.of(new TypeHint<Integer>() {
                            }),
                            TypeInformation.of(new TypeHint<Tuple2<String, String>>() {
                            })
                    );

            /**
             * 对事件流的每条数据进行处理
             * 只能对广播流的数据进行读取操作
             * @param value
             * @param ctx
             * @param out
             * @throws Exception
             */
            @Override
            public void processElement(Integer value, ReadOnlyContext ctx, Collector<Tuple2<String, String>> out) throws Exception {
                //只读操作，意味着只能读取操作，不能修改数据，根据广告id获取广告信息
                ReadOnlyBroadcastState<Integer, Tuple2<String, String>> broadcastState = ctx.getBroadcastState(mapStateDescriptor);
                //根据广告id获取广告位信息
                Tuple2<String, String> tuple2 = broadcastState.get(value);
                if (tuple2 != null) {
                    out.collect(tuple2);
                }
            }

            /**
             * 可以对广播流的数据进行修改
             * @param value
             * @param ctx
             * @param out
             * @throws Exception
             */
            @Override
            public void processBroadcastElement(Map<Integer, Tuple2<String, String>> value,
                                                Context ctx, Collector<Tuple2<String, String>> out) throws Exception {
                //将原来的广播流的数据在状态中删除掉（广播状态意味着可以删除修改数据）
                BroadcastState<Integer, Tuple2<String, String>> broadcastState = ctx.getBroadcastState(mapStateDescriptor);
                //将状态中保存的历史的广播流数据删除掉，避免存在大量垃圾数据
                broadcastState.clear();
                //将最新的广播流数据放入到广播状态中
                broadcastState.putAll(value);
            }
        });

        //TODO 7）打印输出
        adSourceStream.printToErr("广告配置流>>>");
        result.print();

        //TODO 8）运行作业
        env.execute();
    }

    /**
     * 广告位流
     */
    public static class MySourceForBroadcastFunction implements SourceFunction<Map<Integer, Tuple2<String, String>>> {
        private final Random random = new Random();
        private final List<Tuple2<String, String>> ads = Arrays.asList(
                Tuple2.of("baidu", "搜索引擎"),
                Tuple2.of("google", "科技大牛"),
                Tuple2.of("aws", "全球领先的云平台"),
                Tuple2.of("aliyun", "全球领先的云平台"),
                Tuple2.of("腾讯", "氪金使我变强"),
                Tuple2.of("阿里巴巴", "电商龙头"),
                Tuple2.of("字节跳动", "靠算法出名"),
                Tuple2.of("美团", "黄色小公司"),
                Tuple2.of("饿了么", "蓝色小公司"),
                Tuple2.of("瑞幸咖啡", "就是好喝")
        );
        private boolean isRun = true;

        @Override
        public void run(SourceContext<Map<Integer, Tuple2<String, String>>> ctx) throws Exception {
            while (isRun) {
                Map<Integer, Tuple2<String, String>> map = new HashMap<>();
                int keyCounter = 0;
                for (int i = 0; i < ads.size(); i++) {
                    keyCounter++;
                    map.put(keyCounter, ads.get(random.nextInt(ads.size())));
                }
                ctx.collect(map);

                TimeUnit.SECONDS.sleep(5L);
            }
        }

        @Override
        public void cancel() {
            this.isRun = false;
        }
    }
}
